import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
from matplotlib.ticker import MaxNLocator

# =====================
# 图表全局设置
# =====================
plt.rcParams.update({
    'font.size': 12,
    'axes.titlesize': 16,
    'axes.labelsize': 14,
    'xtick.labelsize': 12,
    'ytick.labelsize': 12,
    'legend.fontsize': 12,
    'figure.figsize': (14, 6),
    'font.family': 'DejaVu Sans'  # 避免中文乱码
})

# =====================
# 实验数据准备
# =====================
# 恶意客户端防御数据（表5）
attack_types = ['Gradient Tampering', 'Label Flipping', 'Backdoor Attack']
fedavg_acc = [62.3, 58.1, 67.4]  # 后门攻击为成功率补数
fedatt_acc = [68.5, 64.8, 74.2]
fedpvr_acc = [75.2, 72.6, 81.7]
fedent_acc = [81.7, 79.4, 90.3]  # 后门攻击成功率补数

# 极端Non-IID数据（图5b）
beta_values = [0.05, 0.1, 0.3, 0.5, 1.0]  # Dirichlet参数
fedavg_noniid = [41.2, 48.5, 62.1, 72.3, 82.5]
fedprox_noniid = [48.6, 55.2, 68.7, 75.9, 84.1]
cgofed_noniid = [63.5, 68.9, 76.4, 80.2, 86.7]
fedent_noniid = [72.8, 76.5, 82.1, 85.3, 89.2]

# =====================
# 创建双子图布局
# =====================
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 5))
plt.subplots_adjust(wspace=0.25)

# =====================
# 子图1：恶意客户端防御性能（柱状图）
# =====================
bar_width = 0.2
x_pos = np.arange(len(attack_types))

# 绘制四组柱状对比
patterns = ['/', '\\', 'X', '+']  # 定义四种不同图案
ax1.bar(x_pos - 1.5*bar_width, fedavg_acc, width=bar_width, 
        label='FedAvg', alpha=0.8, color='#1f77b4', hatch=patterns[0])
ax1.bar(x_pos - 0.5*bar_width, fedatt_acc, width=bar_width, 
        label='FedAtt', alpha=0.8, color='#ff7f0e', hatch=patterns[1])
ax1.bar(x_pos + 0.5*bar_width, fedpvr_acc, width=bar_width, 
        label='FedPVR', alpha=0.8, color='#2ca02c', hatch=patterns[2])
ax1.bar(x_pos + 1.5*bar_width, fedent_acc, width=bar_width, 
        label='FedEntGate', alpha=0.8, color='#d62728', hatch=patterns[3])

# 添加数据标签和网格
# for i, acc in enumerate(fedent_acc):
#     ax1.text(i + 1.5*bar_width, acc + 1.5, f'{acc}%', 
#              ha='center', fontsize=11, fontweight='bold')

ax1.set_xticks(x_pos)
ax1.set_xticklabels(attack_types)
ax1.set_ylabel('Accuracy(%)/Defense Success Rate(%)', fontsize=12)
ax1.set_ylim(50, 95)
ax1.grid(axis='y', linestyle='--', alpha=0.7)
ax1.set_title('(a) Model Performance under Malicious Attacks',pad=20, y=-0.35, fontsize=13)

# 添加特殊标注
# ax1.text(2.1, 85, 'Backdoor ASR↓47%', fontsize=11, 
#          bbox=dict(facecolor='gold', alpha=0.3))
ax1.legend(loc='upper left', framealpha=0.9)

# =====================
# 子图2：极端Non-IID鲁棒性（折线图）
# =====================
ax2.plot(beta_values, fedavg_noniid, 'o--', label='FedAvg', 
         linewidth=2, markersize=6, color='#1f77b4')
ax2.plot(beta_values, fedprox_noniid, 's--', label='FedProx', 
         linewidth=2, markersize=6, color='#ff7f0e')
ax2.plot(beta_values, cgofed_noniid, 'D--', label='CGoFed', 
         linewidth=2, markersize=6, color='#2ca02c')
ax2.plot(beta_values, fedent_noniid, 'P-', label='FedEntGate', 
         linewidth=3, markersize=6, color='#d62728')

# 标注关键数据点
# ax2.text(0.05, 72.8+1, 'β=0.05: +14.6%', fontsize=10, 
#          bbox=dict(facecolor='lavender', alpha=0.7))
# ax2.text(0.3, 82.1+1, 'AF=0.03', fontsize=10, color='darkgreen')

# 设置坐标轴和网格
ax2.set_xlabel('Dirichlet Parameter β', fontsize=12)
ax2.set_ylabel('Accuracy (%)', fontsize=12)
ax2.set_title('(b) Model Performance under Extreme Non-IID Scenarios',pad=20, y=-0.35, fontsize=13)
ax2.grid(linestyle='--', alpha=0.6)
ax2.legend(loc='lower right')
ax2.set_ylim(35, 95)

# 添加异构性说明箭头
# ax2.annotate('极端Non-IID', xy=(0.07, 45), xytext=(0.05, 35),
#              arrowprops=dict(arrowstyle="->", color='red'),
#              fontsize=11, color='red')
# ax2.annotate('IID数据', xy=(0.95, 82), xytext=(0.8, 75),
#              arrowprops=dict(arrowstyle="->", color='blue'),
#              fontsize=11, color='blue')

# =====================
# 添加全局标题和保存
# =====================
plt.suptitle('Fig.5 Robustness Verification of FedEntGate: \nMalicious Client Defense & Extreme Non-IID Adaptability', 
             fontsize=18, y=0.99)
plt.tight_layout(rect=[0, 0, 1, 0.95])  # 为全局标题留空间

# 保存高分辨率图片
plt.savefig('Fig5_Robustness_Verification.png', dpi=300, bbox_inches='tight')
plt.show()
